"""
性能监控工具模块
提供Prometheus指标收集和性能监控功能
"""

from prometheus_client import Counter, Gauge, Histogram
import time
from functools import wraps
from typing import Callable, Any
import asyncio

# 定义性能监控指标（使用唯一名称避免重复注册）
REQUEST_COUNT = Counter('agent_requests_total_v2', 'Total agent requests', ['agent_type'])
REQUEST_DURATION = Histogram('agent_request_duration_seconds_v2', 'Agent request duration', ['agent_type'])
ACTIVE_REQUESTS = Gauge('active_requests_v2', 'Currently active requests')
CACHE_HITS = Counter('cache_hits_total_v2', 'Total cache hits')
CACHE_MISSES = Counter('cache_misses_total_v2', 'Total cache misses')
LLM_CALLS = Counter('llm_calls_total_v2', 'Total LLM API calls')
LLM_TOKENS = Counter('llm_tokens_total_v2', 'Total LLM tokens processed')
SEARCH_REQUESTS = Counter('search_requests_total_v2', 'Total search requests')
WORKFLOW_EXECUTIONS = Counter('workflow_executions_total_v2', 'Total workflow executions')
WORKFLOW_DURATION = Histogram('workflow_duration_seconds_v2', 'Workflow execution duration')

def track_performance(agent_type: str):
    """性能监控装饰器"""
    def decorator(func):
        if asyncio.iscoroutinefunction(func):
            @wraps(func)
            async def async_wrapper(*args, **kwargs):
                ACTIVE_REQUESTS.inc()
                start_time = time.time()
                REQUEST_COUNT.labels(agent_type=agent_type).inc()
                
                try:
                    result = await func(*args, **kwargs)
                    duration = time.time() - start_time
                    REQUEST_DURATION.labels(agent_type=agent_type).observe(duration)
                    return result
                finally:
                    ACTIVE_REQUESTS.dec()
            return async_wrapper
        else:
            @wraps(func)
            def sync_wrapper(*args, **kwargs):
                ACTIVE_REQUESTS.inc()
                start_time = time.time()
                REQUEST_COUNT.labels(agent_type=agent_type).inc()
                
                try:
                    result = func(*args, **kwargs)
                    duration = time.time() - start_time
                    REQUEST_DURATION.labels(agent_type=agent_type).observe(duration)
                    return result
                finally:
                    ACTIVE_REQUESTS.dec()
            return sync_wrapper
    return decorator

def track_llm_call():
    """跟踪LLM调用装饰器"""
    def decorator(func):
        if asyncio.iscoroutinefunction(func):
            @wraps(func)
            async def async_wrapper(*args, **kwargs):
                LLM_CALLS.inc()
                return await func(*args, **kwargs)
            return async_wrapper
        else:
            @wraps(func)
            def sync_wrapper(*args, **kwargs):
                LLM_CALLS.inc()
                return func(*args, **kwargs)
            return sync_wrapper
    return decorator

def track_workflow_execution():
    """跟踪工作流执行装饰器"""
    def decorator(func):
        if asyncio.iscoroutinefunction(func):
            @wraps(func)
            async def async_wrapper(*args, **kwargs):
                WORKFLOW_EXECUTIONS.inc()
                start_time = time.time()
                
                try:
                    result = await func(*args, **kwargs)
                    duration = time.time() - start_time
                    WORKFLOW_DURATION.observe(duration)
                    return result
                except Exception as e:
                    duration = time.time() - start_time
                    WORKFLOW_DURATION.observe(duration)
                    raise e
            return async_wrapper
        else:
            @wraps(func)
            def sync_wrapper(*args, **kwargs):
                WORKFLOW_EXECUTIONS.inc()
                start_time = time.time()
                
                try:
                    result = func(*args, **kwargs)
                    duration = time.time() - start_time
                    WORKFLOW_DURATION.observe(duration)
                    return result
                except Exception as e:
                    duration = time.time() - start_time
                    WORKFLOW_DURATION.observe(duration)
                    raise e
            return sync_wrapper
    return decorator

class PerformanceMonitor:
    """性能监控器类"""
    
    def __init__(self):
        self.metrics = {}
        
    def start_timer(self, name: str):
        """开始计时器"""
        self.metrics[name] = {
            'start_time': time.time(),
            'end_time': None,
            'duration': None
        }
    
    def stop_timer(self, name: str) -> float:
        """停止计时器并返回持续时间"""
        if name in self.metrics:
            end_time = time.time()
            duration = end_time - self.metrics[name]['start_time']
            self.metrics[name]['end_time'] = end_time
            self.metrics[name]['duration'] = duration
            return duration
        return 0.0
    
    def get_metrics(self) -> dict:
        """获取所有指标"""
        return self.metrics

# 全局性能监控器实例
performance_monitor = PerformanceMonitor()
